﻿Imports System
Imports System.IO
Imports System.Threading
Imports System.Diagnostics

Public Class ThrottledNetworkStream
    Inherits System.IO.Stream

    Private Shared mHeartBeatThread As System.Threading.Thread
    Private Shared mBytesLastSecondByTrackingId As New Generic.Dictionary(Of String, Integer)
    Private Shared mBytesCurrentSecondByTrackingId As New Generic.Dictionary(Of String, Integer)
    Private mDataBlock As IncomingDataChannelBlock
    Public Sub SetActiveBlock(ByVal dataBlock As IncomingDataChannelBlock)
        mDataBlock = dataBlock
    End Sub
    Public Sub ClearActiveBlock()
        mDataBlock = Nothing
    End Sub

    Public Shared Sub Start()
        If mHeartBeatThread Is Nothing Then
            mHeartBeatThread = New System.Threading.Thread(AddressOf HeartBeatWorker)
            mHeartBeatThread.Name = "Throttled Network Stream Heart Beat Thread"
            mHeartBeatThread.Priority = ThreadPriority.Normal
            mHeartBeatThread.IsBackground = True
            mHeartBeatThread.Start()
        End If
    End Sub
    Public Shared Sub [Stop]()
        If mHeartBeatThread IsNot Nothing Then
            mHeartBeatThread.Abort()
            mHeartBeatThread = Nothing
        End If
    End Sub

    Private Shared mSyncLockSendSetBytes As New Object
    Private Shared mSyncLockReceiveSetBytes As New Object
    Private Shared mSyncLockSendSlotWait As New Object
    Private Shared mSyncLockReceiveSlotWait As New Object

    Private Shared mBytesSent As Integer
    Private Shared mBytesReceived As Integer
    Private Shared mBytesSent_Last As Integer
    Private Shared mBytesReceived_Last As Integer

    Private Shared Sub HeartBeatWorker()
        Try
            Do
                Try
                    SyncLock mSyncLockSendSetBytes
                        mBytesSent_Last = mBytesSent
                        mBytesSent = 0
                    End SyncLock
                    SyncLock mSyncLockReceiveSetBytes
                        mBytesReceived_Last = mBytesReceived
                        mBytesReceived = 0
                    End SyncLock
                    SyncLock mBytesCurrentSecondByTrackingId
                        SyncLock mBytesLastSecondByTrackingId
                            mBytesLastSecondByTrackingId = mBytesCurrentSecondByTrackingId
                            mBytesCurrentSecondByTrackingId = New Generic.Dictionary(Of String, Integer)
                        End SyncLock
                    End SyncLock
                    Thread.Sleep(999)
                Catch ex As ThreadAbortException
                    Exit Do
                Catch ex As Exception
                    Framework.Misc.HandleException(ex)
                End Try
            Loop
        Catch ex As ThreadAbortException
        Catch ex As Exception
            Framework.Misc.HandleException(ex)
        End Try
    End Sub

    Private mBaseStream As Stream

    Public Overloads Overrides ReadOnly Property CanRead() As Boolean
        Get
            Return mBaseStream.CanRead
        End Get
    End Property

    Public Overloads Overrides ReadOnly Property CanSeek() As Boolean
        Get
            Return mBaseStream.CanSeek
        End Get
    End Property

    Public Overloads Overrides ReadOnly Property CanWrite() As Boolean
        Get
            Return mBaseStream.CanWrite
        End Get
    End Property

    Public Overloads Overrides ReadOnly Property Length() As Long
        Get
            Return mBaseStream.Length
        End Get
    End Property

    Public Overloads Overrides Property Position() As Long
        Get
            Return mBaseStream.Position
        End Get
        Set(ByVal value As Long)
            mBaseStream.Position = value
        End Set
    End Property

    Public Sub New(ByVal baseStream As Stream)
        mBaseStream = baseStream
    End Sub

    Public Overloads Overrides Sub Flush()
        mBaseStream.Flush()
    End Sub

    Public Overloads Overrides Function Read(ByVal buffer As Byte(), ByVal offset As Integer, ByVal count As Integer) As Integer
        Throttle(count, False)

        Dim iRet As Integer = mBaseStream.Read(buffer, offset, count)
        If mDataBlock IsNot Nothing AndAlso Not String.IsNullOrEmpty(mDataBlock.TrackingId) Then
            SyncLock mBytesCurrentSecondByTrackingId
                If mBytesCurrentSecondByTrackingId.ContainsKey(mDataBlock.TrackingId) Then
                    mBytesCurrentSecondByTrackingId.Item(mDataBlock.TrackingId) += iRet
                Else
                    mBytesCurrentSecondByTrackingId.Add(mDataBlock.TrackingId, iRet)
                End If
            End SyncLock
        End If
        Return iRet
    End Function

    Public Overloads Overrides Function Seek(ByVal offset As Long, ByVal origin As SeekOrigin) As Long
        Return mBaseStream.Seek(offset, origin)
    End Function

    Public Overloads Overrides Sub SetLength(ByVal value As Long)
        mBaseStream.SetLength(value)
    End Sub

    Public Overloads Overrides Sub Write(ByVal buffer As Byte(), ByVal offset As Integer, ByVal count As Integer)
        Throttle(count, True)

        mBaseStream.Write(buffer, offset, count)
    End Sub

    Public Overloads Overrides Function ToString() As String
        Return mBaseStream.ToString()
    End Function

    ''' <summary> 
    ''' Throttles for the specified buffer size in bytes. 
    ''' </summary> 
    ''' <param name="bufferSizeInBytes">The buffer size in bytes.</param> 
    Protected Sub Throttle(ByVal bufferSizeInBytes As Integer, ByVal isUpstream As Boolean)
        If bufferSizeInBytes = 0 Then Return

        ' Make sure the buffer isn't empty. 
        Dim lMaxBytesPerSecond As Long
        If isUpstream Then
            If Settings.ThrottleUpStream Then
                lMaxBytesPerSecond = Settings.UpStreamBandwidthLimitBytesPerSecond
            Else
                lMaxBytesPerSecond = Long.MaxValue
            End If
            SyncLock mSyncLockSendSlotWait
                While mBytesSent > lMaxBytesPerSecond
                    Thread.Sleep(10)
                End While
            End SyncLock
            SyncLock mSyncLockSendSetBytes
                mBytesSent += bufferSizeInBytes
            End SyncLock
        Else
            If Settings.ThrottleDownStream Then
                lMaxBytesPerSecond = Settings.DownStreamBandwidthLimitBytesPerSecond
            Else
                lMaxBytesPerSecond = Long.MaxValue
            End If
            SyncLock mSyncLockReceiveSlotWait
                While mBytesReceived > lMaxBytesPerSecond
                    Thread.Sleep(10)
                End While
            End SyncLock
            SyncLock mSyncLockReceiveSetBytes
                mBytesReceived += bufferSizeInBytes
            End SyncLock
        End If
    End Sub

    ''' <summary>
    ''' returns rate in bytes per secon
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Shared Function GetCurrentUploadDataRate() As Integer
        Return mBytesSent_Last
    End Function
    ''' <summary>
    ''' returns rate in bytes per secon
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Shared Function GetCurrentDownloadDataRate() As Integer
        Return mBytesReceived_Last
    End Function

    Public Shared Function GetCurrentRatesByTrackingId() As Generic.Dictionary(Of String, Integer)
        Return mBytesLastSecondByTrackingId
    End Function


End Class
